import os
import shutil
import random

def count_photos_in_subfolders(parent_folder):
    subfolders = [f.path for f in os.scandir(parent_folder) if f.is_dir()]
    photo_counts = {}

    for subfolder in subfolders:
        photos = [f for f in os.listdir(subfolder) if f.endswith(('.png', '.jpg', '.jpeg'))]
        photo_counts[subfolder] = len(photos)

    return photo_counts

def move_photos_to_new_folder(parent_folder, new_folder, photo_counts):
    min_photo_count = min(photo_counts.values())

    for subfolder, count in photo_counts.items():
        photos = [f for f in os.listdir(subfolder) if f.endswith(('.png', '.jpg', '.jpeg'))]
        selected_photos = random.sample(photos, min_photo_count)

        new_subfolder_path = os.path.join(new_folder, os.path.basename(subfolder))
        if not os.path.exists(new_subfolder_path):
            os.makedirs(new_subfolder_path)

        for photo in selected_photos:
            shutil.move(os.path.join(subfolder, photo), new_subfolder_path)

parent_folder = 'train_ori'
new_folder = 'train'

photo_counts = count_photos_in_subfolders(parent_folder)
move_photos_to_new_folder(parent_folder, new_folder, photo_counts)
